package main

import (
	"bytes"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"

	_ "embed"
)

var (
	decode    bool
	version   bool
	nopadding bool
)

//go:embed VERSION
var versionStr string

func init() {
	flag.BoolVar(&decode, "d", false, "decode data")
	flag.BoolVar(&decode, "decode", false, "decode data")
	flag.BoolVar(&version, "version", false, "print version")
	flag.BoolVar(&nopadding, "raw", false, "don't use padding")
}

func getInput() (io.ReadCloser, bool, error) {
	args := flag.Args()
	if len(args) == 0 || args[0] == "-" {
		return os.Stdin, true, nil
	}

	f, err := os.Open(args[0])
	if err != nil {
		return nil, false, err
	}
	stat, err := f.Stat()
	if err != nil {
		return nil, false, err
	}
	if stat.IsDir() {
		return nil, false, fmt.Errorf("file is a directory")
	}
	return f, false, nil
}

func getEncoding() *base64.Encoding {
	if nopadding {
		return base64.RawURLEncoding
	} else {
		return base64.URLEncoding
	}
}

func transform() {
	input, isStdin, err := getInput()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return
	}
	defer input.Close()
	encoding := getEncoding()

	if !decode {
		// encoding
		enc := base64.NewEncoder(encoding, os.Stdout)
		defer enc.Close()
		if _, err := io.Copy(enc, input); err != nil {
			fmt.Fprintf(os.Stderr, "\nError encoding: %v", err)
		}
		return
	}

	// decoding
	var r io.Reader
	// 对标准输入的字符串补全
	if !isStdin {
		r = input
	} else {
		// 读取标准输入
		buf := new(bytes.Buffer)
		n, err := buf.ReadFrom(input)
		if err != nil {
			fmt.Fprintf(os.Stderr, "\nError reading from stdin: %v", err)
		}
		if !nopadding {
			if i := int(n % 4); i != 0 {
				buf.WriteString(strings.Repeat("=", 4-i))
			}
		}
		r = buf
	}

	dec := base64.NewDecoder(encoding, r)
	_, err = io.Copy(os.Stdout, dec)
	if err != nil {
		fmt.Fprintf(os.Stderr, "\nError decoding: %v", err)
	}
}

func main() {
	flag.Parse()
	if version {
		fmt.Printf("%s %s\n", os.Args[0], versionStr)
		return
	}
	transform()
	fmt.Println()
}
